import time
import numpy as np
import pandas as pd
import warnings
warnings.filterwarnings("ignore")


def backtest(all_buys):
    stocks_path = r'd:\Python\study_data\stocks\\'
    all_buys = pd.read_csv(all_buys)
    s_5 = []
    s_10 = []
    s_20 = []
    s_30 = []
    s_40 = []
    l_5 = []
    l_10 = []
    l_20 = []
    l_30 = []
    l_40 = []
    cnt = 0
    win_5 = 0
    win_10 = 0
    win_20 = 0
    win_30 = 0
    win_40 = 0
    for i in range(len(all_buys)):
        row = list(all_buys.iloc[i].dropna())
        date = row[0]
        codes = row[1:]
        daily_5 = []
        daily_10 = []
        daily_20 = []
        daily_30 = []
        daily_40 = []
        for code in codes:
            stock_df = pd.read_csv(stocks_path + code + '.csv', index_col=0)
            day = stock_df.index.get_loc(date)
            buy = stock_df.iloc[day]['close']
            if day + 80 < len(stock_df):
                sell_5 = round(stock_df.iloc[day + 5]['close'] / buy - 1, 2)
                sell_10 = round(stock_df.iloc[day + 10]['close'] / buy - 1, 2)
                sell_20 = round(stock_df.iloc[day + 20]['close'] / buy - 1, 2)
                sell_30 = round(stock_df.iloc[day + 40]['close'] / buy - 1, 2)
                sell_40 = round(stock_df.iloc[day + 80]['close'] / buy - 1, 2)
                if np.isinf(sell_5) or np.isinf(sell_10) or np.isinf(sell_20) or np.isinf(sell_30) or np.isinf(sell_40) or np.isnan(sell_5) or np.isnan(sell_10) or np.isnan(sell_20) or np.isnan(sell_30) or np.isnan(sell_40):
                    continue
                else:
                    daily_5.append(sell_5)
                    daily_10.append(sell_10)
                    daily_20.append(sell_20)
                    daily_30.append(sell_30)
                    daily_40.append(sell_40)
        if len(daily_5) > 0:
            cnt += 1
            sell_5_daily = np.mean(daily_5)
            sell_10_daily = np.mean(daily_10)
            sell_20_daily = np.mean(daily_20)
            sell_30_daily = np.mean(daily_30)
            sell_40_daily = np.mean(daily_40)
            if sell_5_daily > 0:
                win_5 += 1
                s_5.append(sell_5_daily)
            else:
                l_5.append(sell_5_daily)
            if sell_10_daily > 0:
                win_10 += 1
                s_10.append(sell_10_daily)
            else:
                l_10.append(sell_10_daily)
            if sell_20_daily > 0:
                win_20 += 1
                s_20.append(sell_20_daily)
            else:
                l_20.append(sell_20_daily)
            if sell_30_daily > 0:
                win_30 += 1
                s_30.append(sell_30_daily)
            else:
                l_30.append(sell_30_daily)
            if sell_40_daily > 0:
                win_40 += 1
                s_40.append(sell_40_daily)
            else:
                l_40.append(sell_40_daily)


            # break
        # break

    pnt = round(win_5 / cnt, 2)
    print(pnt, np.mean(s_5), np.mean(l_5), pnt * np.mean(s_5) + (1 - pnt) * np.mean(l_5))
    pnt = round(win_10 / cnt, 2)
    print(pnt, np.mean(s_10), np.mean(l_10), pnt * np.mean(s_10) + (1 - pnt) * np.mean(l_10))
    pnt = round(win_20 / cnt, 2)
    print(pnt, np.mean(s_20), np.mean(l_20), pnt * np.mean(s_20) + (1 - pnt) * np.mean(l_20))
    pnt = round(win_30 / cnt, 2)
    print(pnt, np.mean(s_30), np.mean(l_30), pnt * np.mean(s_30) + (1 - pnt) * np.mean(l_30))
    pnt = round(win_40 / cnt, 2)
    print(pnt, np.mean(s_40), np.mean(l_40), pnt * np.mean(s_40) + (1 - pnt) * np.mean(l_40))

    return None


if __name__ == '__main__':
    begin_time = time.time()
    test = 'all_buys_qinglong_tupo_backtest.csv'
    backtest(test)

    end_time = time.time()
    run_time = round(end_time - begin_time)
    hour = run_time // 3600
    minute = (run_time - 3600 * hour) // 60
    second = run_time - 3600 * hour - 60 * minute
    print(f'该程序运行时间：{hour}小时{minute}分钟{second}秒')


